import sys
import os
import math
import warnings
import pandas as pd # type: ignore
from scipy.signal import butter, filtfilt, sosfilt, sosfiltfilt, lfilter, firwin, medfilt, wiener, medfilt2d
from scipy.ndimage import median_filter
from scipy import signal
import numpy as np
from statsmodels.stats.outliers_influence import variance_inflation_factor # type: ignore

# models
from sklearn.linear_model import LinearRegression, Ridge, RidgeCV, Lasso, ElasticNet, RANSACRegressor, BayesianRidge # type: ignore
from sklearn.cross_decomposition import PLSRegression, CCA, PLSSVD # type: ignore
from sklearn.kernel_ridge import KernelRidge # type: ignore

import statsmodels.api as sm # type: ignore
import statsmodels.formula.api as smf # type: ignore

from sklearn.metrics import r2_score # type: ignore

def calc_vif(X):
    #X = df[df.columns[13:29]]
    vif = pd.DataFrame()
    vif["Features"] = X.columns
    vif["VIF Factor"] = [variance_inflation_factor(X.values, i) for i in range(X.shape[1])]
    return vif
    print(vif.round(2))

def pd_highpass(input, wn, order=2, type='butter', rp=1., rs=60., _USING_SOS=False):
    sos = None;
    b   = None;
    a   = None;
    
    if _USING_SOS == True:
      if type == 'ellip':
        sos = signal.ellip(order, rp, rs, wn, btype='high', output='sos')
      elif type == 'cheby1':
        sos = signal.cheby1(order, rp, wn, btype='high', output='sos')
      elif type == 'cheby2':
        sos = signal.cheby2(order, rs, wn, btype='high', output='sos')
      elif type == 'bessel':
        sos = signal.bessel(order, wn, btype='high', output='sos')
      else:
        sos = signal.butter(order, wn, btype='high', output='sos')
    else:
      if type == 'ellip':
        b, a = signal.ellip(order, rp, rs, wn, btype='high', output='sos') # type: ignore
      elif type == 'cheby1':
        b, a = signal.cheby1(order, rp, wn, btype='high', output='ba') # type: ignore
      elif type == 'cheby2':
        b, a = signal.cheby2(order, rs, wn, btype='high', output='ba') # type: ignore
      elif type == 'bessel':
        b, a = signal.bessel(order, wn, btype='high', output='ba') # type: ignore
      else:
        b, a = signal.butter(order, wn, btype='high') # type: ignore
    
    out = pd.DataFrame()
    for key in list(input.columns):
      if _USING_SOS == True:
        out[key] = sosfiltfilt(sos, input[key])
      else:
        out[key] = filtfilt(b, a, input[key])
    return out

def pd_lowpass(input, wn, order=2, type='butter', rp=1., rs=60., _USING_SOS=False):
    sos = None;
    b   = None;
    a   = None;
    
    if _USING_SOS == True:
      if type == 'ellip':
        sos = signal.ellip(order, rp, rs, wn, btype='low', output='sos')
      elif type == 'cheby1':
        sos = signal.cheby1(order, rp, wn, btype='low', output='sos')
      elif type == 'cheby2':
        sos = signal.cheby2(order, rs, wn, btype='low', output='sos')
      elif type == 'bessel':
        sos = signal.bessel(order, wn, btype='low', output='sos')
      else:
        sos = signal.butter(order, wn, btype='low', output='sos')
    else:
      if type == 'ellip':
        b, a = signal.ellip(order, rp, rs, wn, btype='low', output='ba') # type: ignore
      elif type == 'cheby1':
        b, a = signal.cheby1(order, rp, wn, btype='low', output='ba') # type: ignore
      elif type == 'cheby2':
        b, a = signal.cheby2(order, rs, wn, btype='low', output='ba') # type: ignore
      elif type == 'bessel':
        b, a = signal.bessel(order, wn, btype='low', output='ba') # type: ignore
      else:
        b, a = signal.butter(order, wn, btype='low', output='ba') # type: ignore
    
    out = pd.DataFrame()
    for key in list(input.columns):
      if _USING_SOS == True:
        out[key] = sosfiltfilt(sos, input[key])
      else:
        out[key] = filtfilt(b, a, input[key])
        # out[key] = filtfilt(b, a, input[key], method='gust')
        
    return out

def pd_bandpass(input, wn, order=2, type='butter', rp=1., rs=60., _USING_SOS=False):
    sos = None;
    b   = None;
    a   = None;
    
    if _USING_SOS == True:
      if type == 'ellip':
        sos = signal.ellip(order, rp, rs, wn, btype='band', output='sos')
      elif type == 'cheby1':
        sos = signal.cheby1(order, rp, wn, btype='band', output='sos')
      elif type == 'cheby2':
        sos = signal.cheby2(order, rs, wn, btype='band', output='sos')
      elif type == 'bessel':
        sos = signal.bessel(order, wn, btype='band', output='sos')
      else:
        sos = signal.butter(order, wn, btype='band', output='sos')
    else:
      if type == 'ellip':
        b, a = signal.ellip(order, rp, rs, wn, btype='band', output='ba') # type: ignore
      elif type == 'cheby1':
        b, a = signal.cheby1(order, rp, wn, btype='band', output='ba') # type: ignore
      elif type == 'cheby2':
        b, a = signal.cheby2(order, rs, wn, btype='band', output='ba') # type: ignore
      elif type == 'bessel':
        b, a = signal.bessel(order, wn, btype='band', output='sos') # type: ignore
      else:
        b, a = signal.butter(order, wn, btype='band', output='ba') # type: ignore
    
    out = pd.DataFrame()
    for key in list(input.columns):
        if _USING_SOS == True:
          out[key] = sosfiltfilt(sos, input[key])
        # elif _USING_DIRECTII == True:
        #   output[key] = filtfilt(b, a, input[key])
        else:
          out[key] = filtfilt(b, a, input[key])
          # out[key] = filtfilt(b, a, input[key], method='gust')
          
    return out
      
def construct_features(input, delta_t=0.1, base_name='basename', output_dir='./'):
    _denominator = (input.iloc[:,0]**2+input.iloc[:,1]**2+input.iloc[:,2]**2)**0.5
    output = pd.DataFrame()
    output['COSX'] = input.iloc[:,0] / _denominator
    output['COSY'] = input.iloc[:,1] / _denominator
    output['COSZ'] = input.iloc[:,2] / _denominator
    
    output.to_csv(output_dir+base_name+'_features.csv', float_format='% 6.6e', sep='\t', index=False)
    
    # 计算导数
    #delta_time = 200. # 1000./200. = 5.
    D_TEMP = pd.DataFrame()
    # D_TEMP['COSX_D'] = (output['COSX'].shift(-1) - output['COSX'].shift(1))*5.
    # D_TEMP['COSY_D'] = (output['COSY'].shift(-1) - output['COSY'].shift(1))*5.
    # D_TEMP['COSZ_D'] = (output['COSZ'].shift(-1) - output['COSZ'].shift(1))*5.
    # D_TEMP = D_TEMP.iloc[1:-1]
    #D_TEMP.to_csv('D_TEMP.csv', index=False)
    D_TEMP['COSX_D'] = np.gradient(output['COSX'], delta_t) # 1/sample_rate=0.1
    D_TEMP['COSY_D'] = np.gradient(output['COSY'], delta_t)
    D_TEMP['COSZ_D'] = np.gradient(output['COSZ'], delta_t)
    D_TEMP.to_csv(output_dir+base_name+'_grad.csv', float_format='% 6.6e', sep='\t', index=False)
    # 构造特征列
    # output = output.iloc[1:-1]
    
    output['COSX*COSX'] = np.multiply(output['COSX'], output['COSX'])
    output['COSX*COSY'] = np.multiply(output['COSX'], output['COSY'])
    output['COSX*COSZ'] = np.multiply(output['COSX'], output['COSZ'])
    output['COSY*COSY'] = np.multiply(output['COSY'], output['COSY'])
    output['COSY*COSZ'] = np.multiply(output['COSY'], output['COSZ'])

    output['COSX*COSX_D'] = np.multiply(output['COSX'], D_TEMP['COSX_D'])
    output['COSX*COSY_D'] = np.multiply(output['COSX'], D_TEMP['COSY_D'])
    output['COSX*COSZ_D'] = np.multiply(output['COSX'], D_TEMP['COSZ_D'])
    output['COSY*COSX_D'] = np.multiply(output['COSY'], D_TEMP['COSX_D'])
    output['COSY*COSY_D'] = np.multiply(output['COSY'], D_TEMP['COSY_D'])
    output['COSY*COSZ_D'] = np.multiply(output['COSY'], D_TEMP['COSZ_D'])
    output['COSZ*COSX_D'] = np.multiply(output['COSZ'], D_TEMP['COSX_D'])
    output['COSZ*COSY_D'] = np.multiply(output['COSZ'], D_TEMP['COSY_D'])

    output['mag_filtered'] = input.iloc[:,3]
    return output

def calc_model(features, model_type='OLS'):
    X_train = features.iloc[:,:-1]
    Y_train = features.iloc[:,-1]
    # X_train.to_csv('X_train.csv', index=False)
    # Y_train.to_csv('Y_train.csv', index=False)
    if model_type == 'OLS':
        model = LinearRegression()
        model.fit(X_train, Y_train)
        coef = model.coef_
    if model_type == 'RidgeCV':
        alphas = np.linspace(0.001, 1, 1000)
        model = RidgeCV(alphas=alphas, cv=5, fit_intercept=False)
        model.fit(X_train, Y_train)
        coef = model.coef_
    elif model_type == 'Ridge':
        model = Ridge(alpha=0.001)
        model.fit(X_train, Y_train)
        coef = model.coef_
    elif model_type == 'PLSR':
        model = PLSRegression(n_components=10)
        model.fit(X_train, Y_train)
        coef = model.coef_.squeeze()
    elif model_type == 'BayesianRidge':
        model = BayesianRidge()
        model.fit(X_train, Y_train)
        coef = model.coef_.squeeze() # type: ignore
    elif model_type == 'SM':
        model = sm.OLS(Y_train, X_train).fit()
        coef = model.params.to_numpy()
        
    # elif model_type == 'CCA':
    #     model = CCA(n_components=10)
    #     model.fit(X_train, Y_train)
    #     coef = model.coef_.squeeze()
    #     # print('select model_type:', model_type)
    # elif model_type == 'PLSSVD':
    #     model = PLSSVD(n_components=10)
    #     model.fit(X_train, Y_train)
    #     coef = model.coef_.squeeze()
    #     # print('select model_type:', model_type)
    # elif model_type == 'Lasso':
    #     model = Lasso(alpha=0.001)
    #     model.fit(X_train, Y_train)
    #     coef = model.coef_.squeeze()
    # elif model_type == 'ElasticNet':
    #     model = ElasticNet()
    #     model.fit(X_train, Y_train)
    #     coef = model.coef_.squeeze()
    # elif model_type == 'RANSACRegressor':
    #     model = RANSACRegressor()
    #     model.fit(X_train, Y_train)
    #     coef = model.coef_
    else: #OLS
        model = LinearRegression()
        model.fit(X_train, Y_train)
        coef = model.coef_
        
        # model = PLSRegression(n_components=10)
        # model.fit(X_train, Y_train)
        # coef = model.coef_.squeeze() 
          
        # model = sm.OLS(Y_train, X_train).fit()
        # coef = model.params
    return model,coef

# calc compensation
def calc_compensation(features, coef):
    X_train = features.iloc[:,:16]
    Y_train = features.iloc[:,-1]
    error = np.matmul(X_train, coef)
    return error
  
# 计算航向角

def calc_angle(filename):
    cwd = os.getcwd()
    input_dir = os.path.join(cwd, 'input/')
    output_dir = os.path.join(cwd, 'output/')
    #intermediate_dir = os.path.join(output_dir + 'intermediate/')

    if os.path.exists('./input/'):
      input_dir = './input/'
      output_dir = './output/'
    elif os.path.exists('../input/'):
      input_dir = '../input/'
      output_dir = '../output/'
    else:
      print('input_dir not found!')
    intermediate_dir = output_dir + 'intermediate/'

    input_file = os.path.join(input_dir, filename)
    print('input', input_file)
    
    df_in = pd.read_csv(input_file, sep=r'\s+') # Read TXT file
    df_out = pd.DataFrame()
    df_out['X分量']      = df_in.iloc[:,0]
    df_out['Y分量']      = df_in.iloc[:,1]
    df_out['Z分量']      = df_in.iloc[:,2]
    df_out['未补磁总场']  = df_in.iloc[:,3]
    df_out['航向角']      = df_out.apply(lambda row: math.atan2(row['Y分量'], row['X分量'])*180/math.pi,axis=1)
    
    #output_file = os.getcwd()/output_name
    basename = os.path.splitext(filename)[0]
    angle_file = intermediate_dir + basename +'_angle.csv'
    df_out.to_csv(angle_file, float_format='%3.2f', sep='\t', index=False)
